package org.loong.crypto.extension.cms.jce;

import java.io.OutputStream;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;

import javax.crypto.SecretKey;

import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.operator.GenericKey;
import org.bouncycastle.operator.OutputEncryptor;
import org.loong.crypto.common.utils.SecureRandomHolder;
import org.loong.crypto.core.algorithm.EncryptionMethod;
import org.loong.crypto.core.algorithm.KeyAlgorithm;
import org.loong.crypto.core.params.KeyGeneratorParameters;
import org.loong.crypto.extension.cms.EnvelopedDataHelper;
import org.loong.crypto.extension.cms.jce.io.CipherOutputStream;
import org.loong.crypto.extension.utils.AlgorithmIdentifierFactory;
import org.loong.crypto.service.CryptoService;

public class JceCMSContentEncryptorBuilder {

    protected EnvelopedDataHelper helper = new EnvelopedDataHelper();

    private final CryptoService cryptoService;

    private final int keySize;

    private SecureRandom random;

    private ASN1ObjectIdentifier encryptionOID;

    public JceCMSContentEncryptorBuilder(CryptoService cryptoService, EncryptionMethod encryptionMethod) throws NoSuchAlgorithmException {
        this(cryptoService, encryptionMethod, encryptionMethod.getKeyLength());
    }

    public JceCMSContentEncryptorBuilder(CryptoService cryptoService, EncryptionMethod encryptionMethod, int keySize) throws NoSuchAlgorithmException {
        this.cryptoService = cryptoService;
        this.encryptionOID = encryptionMethod.getOID();

        int fixedSize = encryptionMethod.getKeyLength();
        if (encryptionOID.equals(PKCSObjectIdentifiers.des_EDE3_CBC)) {
            if (keySize != 168 && keySize != fixedSize) {
                throw new IllegalArgumentException("incorrect keySize for encryptionOID passed to builder.");
            }
            this.keySize = 168;
        } else {
            if (fixedSize > 0 && fixedSize != keySize) {
                throw new IllegalArgumentException("incorrect keySize for encryptionOID passed to builder.");
            }
            this.keySize = keySize;
        }
    }

    public JceCMSContentEncryptorBuilder setSecureRandom(SecureRandom random) {
        this.random = random;
        return this;
    }

    public OutputEncryptor build() throws CMSException {
        return new CMSOutputEncryptor(encryptionOID, keySize, random);
    }

    private class CMSOutputEncryptor implements OutputEncryptor {

        private ASN1ObjectIdentifier encryptionOID;

        private SecretKey encKey;
        
        private byte[] iv;

        private AlgorithmIdentifier algorithmIdentifier;

        CMSOutputEncryptor(ASN1ObjectIdentifier encryptionOID, int keySize, SecureRandom random) throws CMSException {
            this.encryptionOID = encryptionOID;

            String algorithmName = helper.getBaseCipherName(encryptionOID);
            if (algorithmName == null) {
                throw new CMSException("Unsupported encryption algorithm: " + encryptionOID);
            }

            if (random == null) {
                random = SecureRandomHolder.getRandom();
            }
            
            try {
                this.algorithmIdentifier = AlgorithmIdentifierFactory.generateEncryptionAlgID(encryptionOID, random);
                this.encKey = cryptoService.generateKey(KeyAlgorithm.find(algorithmName), KeyGeneratorParameters.builder().keySize(keySize).build());
                this.iv = helper.getIv(encKey, algorithmIdentifier);
            } catch (NoSuchAlgorithmException | IllegalArgumentException e) {
                throw new CMSException(e.getMessage(), e);
            }
        }

        @Override
        public AlgorithmIdentifier getAlgorithmIdentifier() {
            return algorithmIdentifier;
        }

        @Override
        public OutputStream getOutputStream(OutputStream dOut) {
            return new CipherOutputStream(dOut, cryptoService, encryptionOID, encKey, iv);
        }

        @Override
        public GenericKey getKey() {
            return new GenericKey(algorithmIdentifier, encKey.getEncoded());
        }
    }
}
